package com.tuniu.agents.common.agent;

import com.tuniu.agents.advisor.ChatEnvironmentMemory;
import com.tuniu.agents.agent.Agent;
import com.tuniu.agents.agent.AgentManagerGroup;
import com.tuniu.agents.command.Command;
import com.tuniu.agents.command.CommandAction;
import com.tuniu.agents.command.CommandConfig;
import com.tuniu.agents.command.action.CallAgentAction;
import com.tuniu.agents.command.action.CallHandlerAction;
import com.tuniu.agents.command.action.CallToolAction;
import com.tuniu.agents.command.action.ExecuteCustomLogicAction;
import com.tuniu.agents.command.action.OutputTextAction;
import com.tuniu.agents.common.command.CommandConfiguration;
import com.tuniu.agents.common.command.CommandManager;
import com.tuniu.agents.common.command.action.DefaultCallAgentAction;
import com.tuniu.agents.impl.base.Dispatcher;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.model.ChatResponse;
import reactor.core.publisher.Flux;

import java.util.*;
import java.util.stream.Collectors;

@Slf4j
public class CommandAgent extends Agent {

    public static final String AGENT_ID = "CommandAgent";

    /**
     * Global shortcut commands
     */
    private List<Command> globalCommands = new LinkedList<>();

    private CommandManager commandManager;

    private ChatEnvironmentMemory chatEnvironmentMemory;

    private AgentManagerGroup agentManagerGroup;

    public CommandAgent(CommandManager commandManager, AgentManagerGroup agentManagerGroup) {
        this.chatEnvironmentMemory = agentManagerGroup.getChatEnvironmentMemory();
        this.commandManager = commandManager;
        // Add the dispatcher command
        this.globalCommands.add(commandManager.createCommand(
                CommandConfig.builder()
                        .id(Dispatcher.AGENT_ID)
                        .matchName(CommandConfiguration.ALL_MATCH)
                        .commandName("dispatcherAgentCommand")
                        .actionName("defaultCallAgentAction")
                        .priority(1)
                        .params(Map.of(DefaultCallAgentAction.PARAM_AGENT_ID, Dispatcher.AGENT_ID))
                        .build()
        ));
    }

    @Override
    public Flux<String> stream(String conversationId, List<Message> messages, Map<String, Object> param) {
        String input = messages.get(0).getText();
        Map<String, Object> context = chatEnvironmentMemory.get(conversationId);
        // Match shortcut commands based on the conversation, including global and local commands (only used if conditions are met)
        Map<String,Command> partialCommands = commandManager.getAllCommandStates(conversationId);
        TreeMap<String,Command> sortedCommands = new TreeMap<>(partialCommands);
        globalCommands.forEach(
                command -> {
                    // Ignore duplicates
                    if (!sortedCommands.containsKey(command.getId())) {
                        sortedCommands.put(command.getId(), command);
                    }
                }
        );
        List<Command> commandList = sortedCommands.values().stream()
                .sorted(Comparator.comparingInt(Command::getPriority).reversed())
                .collect(Collectors.toList());
        Command command = commandManager.findMatchingCommand(commandList, input, context);
        log.info("execute command: {}", command.getPriority());
        CommandAction action = command.getAction();
        Map<String, Object> concatParam = new HashMap<>(command.getParams());
        if(param != null) {
            concatParam.putAll(param);
        }
        switch (action.getActionType()) {
            case CALL_TOOL:
                CallToolAction toolAction = (CallToolAction)action;
                return toolAction.execute(input, concatParam, context);
            case CALL_AGENT:
                CallAgentAction agentAction = (CallAgentAction)action;
                return agentAction.execute(input, concatParam, context);
            case OUTPUT_TEXT:
                OutputTextAction textAction = (OutputTextAction)action;
                return Flux.just(textAction.execute(input, concatParam, context));
            case EXECUTE_CUSTOM_LOGIC:
                ExecuteCustomLogicAction customLogicAction = (ExecuteCustomLogicAction)action;
                customLogicAction.execute(input, concatParam, context);

            case CALL_HANDLER:
                CallHandlerAction handlerAction = (CallHandlerAction)action;
                return handlerAction.execute(input, concatParam, context);
            default:
                log.warn("unknown action: {}", action.getActionType());
        }
        return null;
    }

    @Override
    public String getAgentId() {
        return AGENT_ID;
    }

    @Override
    public Flux<ChatResponse> streamResponse(String conversationId, List<Message> messages, Map<String, Object> context) {
        return null;
    }

    @Override
    public ChatClient.ChatClientRequestSpec chatClientRequestSpec(String conversationId, List<Message> messages, Map<String, Object> context) {
        return null;
    }

    @Override
    public ChatResponse callResponse(String conversationId, List<Message> messages, Map<String, Object> context) {
        return null;
    }

    @Override
    public String call(String conversationId, List<Message> messages, Map<String, Object> context) {
        return "";
    }
}
